package com.small.xx.ai.spring.advisor.controller;

import com.alibaba.cloud.ai.autoconfigure.memory.RedisChatMemoryAutoConfiguration;
import com.alibaba.cloud.ai.memory.redis.RedisChatMemoryRepository;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

import static org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID;

/**
 * 基于redis的advisor绑定内存记忆案例
 * MessageChatMemoryAdvisor redis
 *
 * @see RedisChatMemoryAutoConfiguration
 */
@RestController
@RequestMapping("/advisor/redis/message")
public class RedisMessageAdvisorController {

    private final int MAX_MESSAGES = 10;
    private final ChatClient.Builder builder;
    private RedisChatMemoryRepository chatMemoryRepository;
    private ChatClient chatClient;
    private MessageWindowChatMemory messageWindowChatMemory;

    public RedisMessageAdvisorController(ChatClient.Builder builder, RedisChatMemoryRepository chatMemoryRepository) {
        this.builder = builder;
        this.chatMemoryRepository = chatMemoryRepository;
        // 设置内存消息最大条数
        messageWindowChatMemory =
            MessageWindowChatMemory.builder().chatMemoryRepository(chatMemoryRepository).maxMessages(MAX_MESSAGES)
                .build();
        this.chatClient =
            builder.defaultAdvisors(MessageChatMemoryAdvisor.builder(messageWindowChatMemory).build()).build();
    }

    @GetMapping("/config")
    public String config(@RequestParam(name = "maxMessage", defaultValue = "10") int maxMessages) {
        // 设置内存消息最大条数
        messageWindowChatMemory =
            MessageWindowChatMemory.builder().chatMemoryRepository(chatMemoryRepository).maxMessages(MAX_MESSAGES)
                .build();
        this.chatClient =
            builder.defaultAdvisors(MessageChatMemoryAdvisor.builder(messageWindowChatMemory).build()).build();
        return "改为：" + maxMessages;
    }

    @GetMapping("/call")
    public String call(@RequestParam(value = "query", defaultValue = "你好，我叫xx,的外号是XX，请记住呀") String query,
        @RequestParam(value = "conversation_id", defaultValue = "XX") String conversationId) {
        String content =
            chatClient.prompt(query).advisors(a -> a.param(CONVERSATION_ID, conversationId)).call().content();
        System.out.println(content);
        return content;
    }

    @GetMapping("/messages")
    public List<Message> messages(@RequestParam(value = "conversation_id", defaultValue = "XX") String conversationId) {
        return messageWindowChatMemory.get(conversationId);
    }

}
